{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import lsqfit\n",
    "from model_avg_paper import *\n",
    "from model_avg_paper.test_model_vary import test_vary_poly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p0_test_quad = {\n",
    "    'a0': 1.80,\n",
    "    'a1': -0.53,\n",
    "    'a2': 0.31,\n",
    "}\n",
    "Nt = 16\n",
    "noise_params = {\n",
    "    'noise_amp': 1.0,\n",
    "    'noise_samples': 160,\n",
    "    'frac_noise': False,\n",
    "}\n",
    "obs_name='a0'\n",
    "\n",
    "def quad_model(x,p):\n",
    "    return poly_model(x,p,Nt,m=2)\n",
    "\n",
    "# Set seed for consistency of outcome\n",
    "np.random.seed(96737)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = gen_synth_data(\n",
    "    np.arange(1,Nt),\n",
    "    p0_test_quad, \n",
    "    quad_model,\n",
    "    **noise_params)\n",
    "\n",
    "test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_res = test_vary_poly(test_data, m_max=5, m_min=0, Nt=Nt, obs_name=obs_name, IC='AIC')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare table for paper\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "test_res_df = pd.DataFrame({ k: test_res[k] for k in ['obs', 'Qs', 'ICs', 'fits', 'probs']})\n",
    "for pn in ('a1', 'a2', 'a3', 'a4', 'a5'):\n",
    "    test_res_df[pn] = test_res_df['fits'].apply(lambda f: f.p.get(pn))\n",
    "test_res_df['chi2'] = test_res_df['fits'].apply(lambda f: f.chi2)\n",
    "test_res_df['probs'] /= np.sum(test_res_df['probs'])\n",
    "tdf = test_res_df.drop('fits', axis=1).rename(columns={'obs': 'a0'})\n",
    "tdf['probs'] = tdf['probs'].round(2)\n",
    "tdf['chi2'] = tdf['chi2'].round(2)\n",
    "tdf['Qs'] = tdf['Qs'].round(2)\n",
    "tdf['ICs'] = -2*tdf['ICs'].round(2)\n",
    "\n",
    "print(test_res['obs_avg'])\n",
    "\n",
    "tdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Figure 1A\n",
    "\n",
    "plot_gvcorr(test_res['data']['y'], x=np.arange(1,Nt), color='forestgreen')\n",
    "\n",
    "TF = test_res['fits'][2]\n",
    "plot_gvcorr(TF.fcn(np.arange(0,18), p0_test_quad), \n",
    "            x=np.arange(0,18), color='k', fill=True)\n",
    "\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y(x)$')\n",
    "plt.xlim(0,17)\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "# Uncomment to save figure to disk\n",
    "#plt.savefig('plots/poly_data.pdf', bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Figure 1B\n",
    "\n",
    "import matplotlib.ticker as ticker\n",
    "\n",
    "gs = plt.GridSpec(2, 1, height_ratios=[2,1])\n",
    "gs.update(hspace=0.06)\n",
    "\n",
    "ax1 = plt.subplot(gs[0])\n",
    "\n",
    "plot_gvcorr([test_res['obs_avg']], x=np.array([-0.1]), color='red', markersize=9, marker='s', open_symbol=True, label='Model avg.')\n",
    "plot_gvcorr(test_res['obs'], x=test_res['m'], label='Individual fits')\n",
    "\n",
    "ax1.plot(np.array([-1,10]), 0*np.array([0,0])+p0_test_quad[obs_name], linestyle='--', color='k', label='Model truth')\n",
    "ax1.set_xlabel('$m$')\n",
    "ax1.set_ylabel('$a_0$')\n",
    "\n",
    "ax1.legend(loc='center left', bbox_to_anchor=(1,0.5))\n",
    "ax1.set_xlim(-0.3, 5.3)\n",
    "\n",
    "plt.setp(ax1.get_xticklabels(), visible=False)\n",
    "\n",
    "ax2 = plt.subplot(gs[1])\n",
    "ax2.set_xlim(-0.3, 5.3)\n",
    "ax2.set_ylim(0,0.65)\n",
    "\n",
    "p_norm = test_res['probs'] / np.sum(test_res['probs'])\n",
    "Q_norm = test_res['Qs'] / np.sum(test_res['Qs'])\n",
    "plt.plot(test_res['m'], p_norm, color='orange', label='pr$(M|D)$')\n",
    "plt.plot(test_res['m'], Q_norm, color='blue', linestyle='-.', label='Fit p-value')  # Note: fit prob != model prob! \n",
    "\n",
    "tick_spacing = 1\n",
    "ax2.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))\n",
    "plt.yticks([0,np.max(p_norm)])\n",
    "ax2.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: '0' if x == 0 else '{:.2f}'.format(x)))\n",
    "\n",
    "\n",
    "\n",
    "# Put a legend to the right of the current axis\n",
    "ax2.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "ax2.set_xlabel('$m$')\n",
    "ax2.set_ylabel('p')\n",
    "\n",
    "# Uncomment to save figure to disk\n",
    "#plt.savefig('plots/poly_avg.pdf', bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scaling w/number of samples\n",
    "\n",
    "Nsamp_array = np.array([20, 40, 80, 160, 320, 640, 2040, 4096, 4096*2, 4096*4])\n",
    "Nsamp_max = Nsamp_array[-1]\n",
    "\n",
    "noise_params['noise_samples'] = Nsamp_max\n",
    "scale_data = gen_synth_data(\n",
    "    np.arange(1,Nt), \n",
    "    p0_test_quad, \n",
    "    quad_model,\n",
    "    **noise_params)\n",
    "\n",
    "model_avg_vs_Nsamp = []\n",
    "fixed_quad_vs_Nsamp = []\n",
    "fixed_quart_vs_Nsamp = []\n",
    "naive_avg_vs_Nsamp = []\n",
    "fw_vs_Nsamp = []\n",
    "\n",
    "for Nsamp in Nsamp_array:\n",
    "    test_data_scale = cut_synth_data_Nsamp(scale_data, Nsamp)\n",
    "    test_res_scale = test_vary_poly(test_data_scale, m_max=5, m_min=0, Nt=Nt, obs_name=obs_name, IC='AIC')\n",
    "    test_res_scale_naive = test_vary_poly(test_data_scale, m_max=5, m_min=0, Nt=Nt, obs_name=obs_name, IC='naive')\n",
    "    \n",
    "    model_avg_vs_Nsamp.append(test_res_scale['obs_avg'])\n",
    "    naive_avg_vs_Nsamp.append(test_res_scale_naive['obs_avg'])\n",
    "    fixed_quad_vs_Nsamp.append(test_res_scale['obs'][2])\n",
    "    fixed_quart_vs_Nsamp.append(test_res_scale['obs'][4])\n",
    "    fw_vs_Nsamp.append(obs_avg_full_width(test_res_scale['obs'], test_res_scale['Qs'], test_res_scale['fits'], bf_i=2))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Figure 2\n",
    "\n",
    "plot_gvcorr(model_avg_vs_Nsamp, x=np.log(Nsamp_array)+0.1, marker='o', markersize=6, label='Model avg. (AIC)')\n",
    "plot_gvcorr(fixed_quad_vs_Nsamp, x=np.log(Nsamp_array)+0.2, color='red', marker='s', markersize=6,label='Quadratic fixed')\n",
    "plot_gvcorr(fw_vs_Nsamp, x=np.log(Nsamp_array)+0.3, color='orange', marker='X', markersize=8,label='Full-width systematic')\n",
    "plot_gvcorr(naive_avg_vs_Nsamp, x=np.log(Nsamp_array)+0.4, color='silver', marker='v', markersize=8,label='Model avg. (naive)')\n",
    "\n",
    "plt.plot(np.array([0,10]), 0*np.array([0,10])+p0_test_quad[obs_name], linestyle='--', color='k', label='Model truth')\n",
    "plt.xlabel('$\\\\log(N)$')\n",
    "plt.ylabel('$a_0$')\n",
    "plt.xlim(2.7,7.)\n",
    "plt.legend(bbox_to_anchor=(1.01, 1), loc='upper left')\n",
    "\n",
    "# Put a legend to the right of the current axis\n",
    "ax = plt.subplot(111)\n",
    "ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "\n",
    "#plt.show()\n",
    "#plt.ylim(0.78,0.82)\n",
    "#plt.tight_layout()\n",
    "\n",
    "# Uncomment to save figure to disk\n",
    "#plt.savefig('plots/poly_N_scaling.pdf', bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
